import torch.nn as nn
import torch.nn.functional as F

class SegmentLoss(nn.Module):
    def __init__(self, reduction='mean'):
        super().__init__()
        self.reduction = reduction

    def forward(self, output, target, weight=None):
        # input: (n, c, h, w), target: (n, h, w)
        n, c, h, w = output.size()
        log_p = F.log_softmax(input, dim=1)
        # log_p: (n*h*w, c)
        log_p = log_p.permute(0, 2, 3, 1).contiguous()
        log_p = log_p[target.view(n, h, w, 1).repeat(1, 1, 1, c) >= 0]
        log_p = log_p.view(-1, c)
        # target: (n*h*w,)
        mask = target >= 0
        target = target[mask]
        loss = F.nll_loss(log_p, target, weight=weight, reduction=self.reduction)
        return loss 
    